import tensorflow as tf
from tensorflow.python import keras
import matplotlib as mpl
import matplotlib.pyplot as plt
import numpy as np
import sklearn
import pandas as pd
import sys
import os
import time


def print_version():
    print(tf.__version__)
    print(sys.version_info)
    for module in tf, keras, mpl, np, pd, sklearn:
        print(module.__name__, module.__version__)


def load_datasets():
    # 数据集
    fashion_mnist = keras.datasets.fashion_mnist
    (train_images, train_labels), (test_images, test_labels) = fashion_mnist.load_data()
    valid_images, train_images = train_images[:5000], train_images[5000:]
    valid_labels, train_labels = train_labels[:5000], train_labels[5000:]
    print("train datasets:", train_images.shape, train_labels.shape)
    print("valid datasets:", valid_images.shape, valid_labels.shape)
    print("test datasets:", test_images.shape, test_labels.shape)
    return train_images, train_labels, valid_images, valid_labels, test_images, test_labels


def show_image(n_rows, n_cols, x_data, y_data, class_names):
    assert len(x_data) == len(y_data)
    assert n_rows * n_cols < len(x_data)
    plt.figure(figsize=(n_cols * 1.4, n_rows * 1.6))
    for row in range(n_rows):
        for col in range(n_cols):
            index = n_cols * row + col
            plt.subplot(n_rows, n_cols, index + 1)
            plt.imshow(x_data[index], cmap="binary", interpolation="nearest")
            plt.axis("off")
            plt.title(class_names[y_data[index]])
    plt.show()


def create_model():
    model = keras.models.Sequential()
    model.add(keras.layers.Flatten(input_shape=[28, 28]))
    model.add(keras.layers.Dense(300, activation="relu"))
    model.add(keras.layers.Dense(100, activation="relu"))
    model.add(keras.layers.Dense(10, activation="softmax"))
    model.compile(
        loss="sparse_categorical_crossentropy",
        optimizer="sgd",
        metrics=["accuracy"]
    )
    # 摘要
    model.summary()
    return model


if __name__ == '__main__':
    # 打印版本
    # print_version()

    # 加载数据集
    train_images, train_labels, valid_images, valid_labels, test_images, test_labels = load_datasets()

    # 显示数据集
    # class_names = ['T-shirt/top', 'Trouser', 'Pullover', 'Dress', 'Coat', 'Sandal', 'Shirt', 'Sneaker', 'Bag',
    #                'Ankle boot']
    # show_image(10, 10, train_images, train_labels, class_names)

    # 创建模型
    model = create_model()

    # 训练模型
    hs = model.fit(train_images, train_labels, validation_data=(valid_images, valid_labels), nb_epoch=5)

